# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: xguo

import logging
import ujson as json
import math

from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.bitflyer_v1.book_builder import BookBuilder
from coin.exchange.bitflyer_v1.util import get_topic_string
from coin.flow.topic_map import get_default_topic_map


# Return None is no valid book date can be found.
def iter_book_data(raw_message):
  payload_message = raw_message.get('M')
  if not payload_message:
    return

  for msg_elem in payload_message:
    assert msg_elem['H'] == 'BFEXHub', msg_elem
    if msg_elem['M'] != 'ReceiveTickers':
      continue
    assert len(msg_elem['A']) == 1, len(msg_elem['A'])
    for book_msg in msg_elem['A'][0]:
      yield book_msg


class BitflyerHandler(WebsocketHandler):
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._topic_map = {}

  def _build_topic_map(self, native_symbol, product):
    if product is None:
      logging.info('Undefined product: %s' % product)
      self._topic_map[native_symbol] = set()
      return

    try:
      topics = get_topic_string(product, worker_id='*')
    except ValueError:
      logging.info('No subscription: %s' % product)
      self._topic_map[native_symbol] = set()
      return

    topic_id_set = set()
    for topic in topics:
      for topic_str, topic_id in get_default_topic_map().search_topic_id(topic):
        topic_id_set.add(topic_id)
    self._topic_map[native_symbol] = topic_id_set

  def _is_feed_from_legitimate_topic(self, native_symbol, topic_id, current_ts):
    if native_symbol not in self._topic_map:
      product = BitflyerFuturesProduct.FromStrNativeProductNothrow(native_symbol, current_ts)
      self._build_topic_map(native_symbol, product)
    return topic_id in self._topic_map[native_symbol]

  def on_message(self, record, queue_data, topic_data):
    raw_message = record.data
    if not isinstance(raw_message, dict):
      raw_message = json.loads(raw_message.tobytes())

    native_symbol = None
    updated_books = {}
    for message in iter_book_data(raw_message):
      native_symbol = message['pc']  # product code
      if ((self._is_live is False) and not self._is_feed_from_legitimate_topic(
          native_symbol, record.topic_id, record.timestamp)):
        # Drop the message.
        continue

      if not self.is_subscribed("book", native_symbol):
        continue

      # TODO(xguo): Not sure what is ['tk']['ts'], but these types seems
      # to cause problems.
      if message['tk']['ts'] in ('PREOPEN', 'CLOSED', 'SHUTDOWN', 'SPECIAL'):
        continue
      assert message['tk']['ts'] == 'OPEN'
      if message['tt'] == 'UPDATE':
        book_builder = self._update_book(record.timestamp, native_symbol, message, False)
        if book_builder:
          updated_books[native_symbol] = book_builder

        trade_event = self._update_trade(record.timestamp,
                                         native_symbol,
                                         message,
                                         topic_id=record.topic_id)
        if trade_event:
          self.publish_event('trade', native_symbol, trade_event)

      elif message['tt'] == 'SNAPSHOT':
        book_builder = self._update_book(record.timestamp, native_symbol, message, True)
        if book_builder:
          updated_books[native_symbol] = book_builder

      else:
        raise ValueError('Unknown ticker type! %s' % message)

    for native_symbol, book_builder in updated_books.items():
      if self._feed_checker:
        self._feed_checker.on_book(product=book_builder.product,
                                   book=book_builder,
                                   timestamp=record.timestamp)
      book_builder.update_topic_id(record.topic_id)
      self.publish_event('book', native_symbol, book_builder)

  def _update_book(self, timestamp, native_symbol, message, snapshot):
    if not snapshot:
      if native_symbol not in self._books:
        return
      book_builder = self._books[native_symbol]

    else:  # Snapshot
      if native_symbol not in self._books:
        book_builder = BookBuilder(native_symbol, book_builder_observer=self._book_builder_observer)
        self._books[native_symbol] = book_builder
      else:
        book_builder = self._books[native_symbol]
        book_builder.clear_book()

    book_builder.update(message['bd']['b'], message['bd']['a'])
    book_builder.update_timestamp(timestamp)

    ticker_bid0 = float(message['tk']['bb'])
    ticker_ask0 = float(message['tk']['ba'])

    if (book_builder.bid0() is not None
        and not math.isclose(book_builder.bid0().price, ticker_bid0)):
      book_builder._delete_bid_up_to(ticker_bid0)
    if (book_builder.ask0() is not None
        and not math.isclose(book_builder.ask0().price, ticker_ask0)):
      book_builder._delete_ask_up_to(ticker_ask0)

    if book_builder.ask0() is not None and book_builder.bid0() is not None:
      book_ask0 = book_builder.ask0().price
      book_bid0 = book_builder.bid0().price
      assert (book_bid0 < book_ask0)
      assert (book_bid0 <= ticker_bid0)
      assert (book_ask0 >= ticker_ask0)

    return book_builder

  def _update_trade(self, timestamp, native_symbol, message, **kwargs):
    data = message['tk']
    qty = float(data['tv'])
    if qty < 0 or math.isclose(qty, 0.):
      return None

    price = float(data['LTP'])
    if data['tis'] == "BUY":
      side = TradeSide.TRADE_BUY_SIDE
    elif data['tis'] == "SELL":
      side = TradeSide.TRADE_SELL_SIDE
    else:
      side = TradeSide.TRADE_NO_SIDE

    trade_event = create_market_trade_event(timestamp, native_symbol, price, qty, side, **kwargs)
    if self._feed_checker:
      product = BitflyerFuturesProduct.FromStrNativeProduct(native_symbol)
      self._feed_checker.on_trade(product=product, trade_event=trade_event, timestamp=timestamp)
    return trade_event
